from camel.toolkits import SearchToolkit, MathToolkit
from camel.types.agents import ToolCallingRecord
from camel.models import ModelFactory
from camel.types import ModelPlatformType
from camel.societies import RolePlaying
from camel.utils import print_text_animated
from colorama import Fore
import os
from dotenv import load_dotenv

load_dotenv()
api_key = os.getenv('MODELSCOPE_SDK_TOKEN')

def setup_tools():
    """设置工具包"""
    tools_list = [
        *SearchToolkit().get_tools(),
        *MathToolkit().get_tools()
    ]
    return tools_list

def setup_task():
    """设置任务提示"""
    task_prompt = (
        "假设现在是2024年，"
        "搜索牛津大学的成立年份，并计算出其当前年龄。"
        "然后再将这个年龄加上10年。使用百度搜索工具。"
    )
    return task_prompt

def create_model():
    """创建模型实例"""
    model = ModelFactory.create(
        model_platform=ModelPlatformType.OPENAI_COMPATIBLE_MODEL,
        model_type="Qwen/Qwen2.5-72B-Instruct",
        url='https://api-inference.modelscope.cn/v1/',
        api_key=api_key
    )
    return model

def setup_role_playing(tools_list, task_prompt, model):
    """设置角色扮演会话"""
    role_play_session = RolePlaying(
        assistant_role_name="搜索者",
        user_role_name="教授",
        assistant_agent_kwargs=dict(
            model=model,
            tools=tools_list,
        ),
        user_agent_kwargs=dict(
            model=model,
        ),
        task_prompt=task_prompt,
        with_task_specify=False,
        output_language='中文'
    )
    return role_play_session

def print_session_info(role_play_session, task_prompt):
    """打印会话信息"""
    print(
        Fore.GREEN
        + f"AI助手系统消息:\n{role_play_session.assistant_sys_msg}\n"
    )
    print(
        Fore.BLUE + f"AI用户系统消息:\n{role_play_session.user_sys_msg}\n"
    )

    print(Fore.YELLOW + f"原始任务提示:\n{task_prompt}\n")
    print(
        Fore.CYAN
        + "指定的任务提示:"
        + f"\n{role_play_session.specified_task_prompt}\n"
    )
    print(Fore.RED + f"最终任务提示:\n{role_play_session.task_prompt}\n")

def run_conversation(role_play_session, chat_turn_limit=10):
    """运行对话"""
    n = 0
    input_msg = role_play_session.init_chat()

    while n < chat_turn_limit:
        n += 1
        assistant_response, user_response = role_play_session.step(input_msg)

        if assistant_response.terminated:
            print(
                Fore.GREEN
                + (
                    "AI助手终止。原因: "
                    f"{assistant_response.info['termination_reasons']}."
                )
            )
            break
        if user_response.terminated:
            print(
                Fore.GREEN
                + (
                    "AI用户终止。"
                    f"原因: {user_response.info['termination_reasons']}."
                )
            )
            break

        # 打印用户的输出
        print_text_animated(
            Fore.BLUE + f"AI用户:\n\n{user_response.msg.content}\n"
        )

        if "CAMEL_TASK_DONE" in user_response.msg.content:
            break

        # 打印助手的输出，包括任何函数执行信息
        print_text_animated(Fore.GREEN + "AI助手:")
        tool_calls: list[ToolCallingRecord] = assistant_response.info[
            'tool_calls'
        ]
        for func_record in tool_calls:
            print_text_animated(f"{func_record}")
        print_text_animated(f"{assistant_response.msg.content}\n")

        input_msg = assistant_response.msg

def main():
    """主函数"""
    # 设置工具和任务
    tools_list = setup_tools()
    task_prompt = setup_task()
    
    # 创建模型和会话
    model = create_model()
    role_play_session = setup_role_playing(tools_list, task_prompt, model)
    
    # 打印会话信息
    print_session_info(role_play_session, task_prompt)
    
    # 运行对话
    run_conversation(role_play_session)

if __name__ == "__main__":
    main()